#!/usr/bin/env python
import acado

import rospy
import tf2_ros
from geometry_msgs.msg import PoseStamped, TransformStamped,PoseArray,Pose,PoseWithCovarianceStamped
import numpy as np
from ros_numpy import numpify,msgify
from nav_msgs.msg import Path
from sensor_msgs.msg import PointCloud2
from tf.transformations import quaternion_from_euler,euler_matrix
from sensor_msgs.msg import JointState


N=20 #timestep
NU=4
NOD=3
NX=16
NY=16
NYN=4

NW=16
NWN=12

def calculate_observation_cov(T):
    R=T[:3,:3]
    t=T[:3,3]
    d=np.sqrt(np.sum(t**2))
    d_min=0.3
    V=np.zeros([3,3])
    V[0,0]=(d-d_min)**2*100.0
    V[1,1]=(d-d_min)**2*1.0
    V[2,2]=(d-d_min)**2*1.0
    
    return V

def calculate_visual_limit(center_vector,fov_pitch,fov_yaw):
    roll_c,pitch_c,yaw_c=euler_from_vector(center_vector)
    
    pitch_min=pitch_c-fov_pitch/2
    pitch_max=pitch_c+fov_pitch/2
    yaw_min=yaw_c-fov_yaw/2
    yaw_max=yaw_c+fov_yaw/2
    
    min_angles=np.array([pitch_min,yaw_min])
    max_angles=np.array([pitch_max,yaw_max])
    
    pitch_limit=np.array([pitch_min,pitch_max])
    yaw_limit=np.array([yaw_min,yaw_max])
    
    return min_angles,max_angles


def transform_vector(matrix, vector):
    # Extend the 3D vector with a fourth component
    extended_vector = np.append(vector, 1)

    # Perform the transformation by multiplying the matrix and the extended vector
    transformed_vector = np.matmul(matrix, extended_vector)

    # Normalize the result by dividing the first three components by the fourth component
    normalized_vector = transformed_vector[:3] / transformed_vector[3]

    return normalized_vector


def euler_from_vector(vector):
    """Calculate Euler angles from a unit vector representing the new x-axis."""
    # Normalize the input vector
    vector = vector / np.linalg.norm(vector)
    x, y, z = vector

    # Calculate yaw (around z-axis)
    yaw = np.arctan2(y, x)

    # Calculate pitch (around new y-axis)
    pitch = np.arctan2(-z, np.sqrt(x**2 + y**2))

    # Calculate roll (around new x-axis)
    roll = 0.0  # Assuming roll is zero as it depends on the original coordinate system

    return roll, pitch, yaw


def pc2_msg_to_np(pc_msg):
    pc_np=numpify(pc_msg)
    pc_np= np.stack([pc_np[f] for f in ('x', 'y', 'z')])

    return pc_np

class TransformListenerNode:
    def __init__(self):
        # Initialize the node
        rospy.init_node('transform_listener')

        # Create a TransformListener
        self.tf_buffer = tf2_ros.Buffer()
        self.tf_listener = tf2_ros.TransformListener(self.tf_buffer)

        # Create a publisher for the body pose
        self.pose_pub = rospy.Publisher('body_pose_topic', PoseStamped, queue_size=10)
        self.pose_pub1 = rospy.Publisher('body_pose_topic1', PoseStamped, queue_size=10)
        self.path_pub = rospy.Publisher('/slot/path', Path, queue_size=10)
        self.pose_arr_pub=rospy.Publisher('/slot/pose_array',PoseArray,queue_size=10)
        # Start listening to the transform
        
        self.pose_with_cov_pub=rospy.Publisher('/slot/target/pose_with_cov',PoseWithCovarianceStamped,queue_size=10)
        self.observation_cov_pub=rospy.Publisher('/slot/target/observation_cov',PoseWithCovarianceStamped,queue_size=10)
        
        rospy.Subscriber('/slot/pointcloud2',PointCloud2,self.pc_cb)
        
        self.v=np.zeros(3)
        self.T_T_V=np.eye(4)
        self.T_B_V=np.eye(4)
        self.T_T_B=np.eye(4)
        self.pc_B=np.zeros([3,100])
        
        self.kf_P=np.diag([1,1,1])
        
        self.listen_transform()
        
        
        
        
        
    def pc_cb(self,pc_msg):
        
        self.pc_B=pc2_msg_to_np(pc_msg)
        
        

    def listen_transform(self):
        rate = rospy.Rate(10)  # 10 Hz
        
        

        while not rospy.is_shutdown():
            try:
                # Lookup the transform from 'vrpn' to 'target'
                transform = self.tf_buffer.lookup_transform( 'vrpn','target', rospy.Time(0), rospy.Duration(1.0))
                
                self.T_T_V=numpify(transform.transform)
                


                # Lookup the transform from 'vrpn' to 'body'
                transform = self.tf_buffer.lookup_transform( 'vrpn','body',rospy.Time(0), rospy.Duration(1.0))
                                # Create a pose message and set its header frame ID as 'body'
                
                self.T_B_V=numpify(transform.transform)
                
                # Lookup the transform from 'vrpn' to 'body'
                transform = self.tf_buffer.lookup_transform( 'body','target',rospy.Time(0), rospy.Duration(1.0))
                                # Create a pose message and set its header frame ID as 'body'
                
                self.T_T_B=numpify(transform.transform)
                
                
                t_T_B=self.T_T_B[:3,3]
                dist_T_B=np.sqrt(np.sum(t_T_B**2))
                print("dist: "+str(dist_T_B))
                
                
                #publish target pose with cov
                pose_with_cov=PoseWithCovarianceStamped()
                cov_arr=np.zeros([6,6])
                cov_arr[:3,:3]=self.kf_P
                pose_with_cov.pose.covariance=cov_arr.flatten().tolist()
                pose_with_cov.header.frame_id='target'
                self.pose_with_cov_pub.publish(pose_with_cov)
                
                #publish observation cov
                observation_cov_msg=PoseWithCovarianceStamped()
                V=calculate_observation_cov(self.T_T_B)
                
                rpy=euler_from_vector((t_T_B))
                T=euler_matrix(rpy[0],rpy[1],rpy[2])
                R=T[:3,:3]
                S=np.zeros([6,6])
                S[:3,:3]=R @ V @ R.T
                T[:3,3]=t_T_B
                
                observation_cov_msg.pose.pose=msgify(Pose,T)
                observation_cov_msg.pose.covariance=S.flatten().tolist()
                observation_cov_msg.header.frame_id='body'
                self.observation_cov_pub.publish(observation_cov_msg)
                
                
                
                pc_center_B=np.mean(self.pc_B,axis=1)
                tg_center_B=self.T_T_B[:3,3]
                
                pc_weight=0.5313*0
                tg_weight=1-pc_weight
                
                
                pc_weight=0.2313
                tg_weight=1-pc_weight
                
                v_target_B=self.T_T_B[:3,3]
                v_pccent_B=np.mean(self.pc_B,axis=1)
                
                v_desire_B=pc_weight*v_pccent_B+tg_weight*v_target_B
                v_desire_V=transform_vector(self.T_B_V,v_desire_B)
                v_target_V=transform_vector(self.T_B_V,v_target_B)

                
                
                v_B=pc_center_B*pc_weight+tg_center_B*tg_weight
                
                v_V=transform_vector(self.T_B_V,v_B)
                
                #print(np.mean(self.pc_B,axis=1))
                
                x0=np.zeros(NX).reshape(1,NX)
                u=np.zeros([N,NU])

                Q=np.eye(NW)
                
                #for ellipse
                Q=np.diag([0,0,0,1,1,1,1,600,60,60,0,1000,1,1,1,5e5])
                #Q=np.diag([60,60,60,1,1,1,1,600,60,60,1000,1000,10,10,10,5e5])
                
                #straight line
                #Q=np.diag([0,0,0,1,1,1,1,600,60,60,160,1000,1,1,1,5e5])
                
                #static
                #Q=np.diag([0,0,0,1,1,1,1,60,60,60,160,1000,1,1,1,5e5])
                
                W=np.transpose(np.tile(Q,N))
                WN=Q[:-4,:-4]
                y=np.zeros([N,NY])
                
                OD=np.ones([N+1,NOD])
                
                
                x0[0,0]=self.T_B_V[0,3] #px
                x0[0,1]=self.T_B_V[1,3] #py
                x0[0,2]=self.T_B_V[2,3] #pz
                x0[0,3]=1.0#qw
                x0[0,4]=0.0#qx
                x0[0,5]=0.0#qy
                x0[0,6]=0.0#qz
                x0[0,7]=self.v[0]#vx
                x0[0,8]=self.v[1]#vy
                x0[0,9]=self.v[2]#vz
                x0[0,10]=self.kf_P[0,0]#pxx
                x0[0,11]=self.kf_P[0,1]#pxy
                x0[0,12]=self.kf_P[0,2]#pxz
                x0[0,13]=self.kf_P[1,1]#pyy
                x0[0,14]=self.kf_P[1,2]#pyz
                x0[0,15]=self.kf_P[2,2]#pzz
                
                
                
                OD[:,0]=self.T_T_V[0,3]
                OD[:,1]=self.T_T_V[1,3]
                OD[:,2]=self.T_T_V[2,3]
                
                
                X=np.tile(x0.reshape(NX,1),N+1).T
                
                X,u=acado.mpc(1,1,x0,X,u,OD,y,W,WN,1)
                x0=X[1,:].reshape(1,NX)
                
                P=np.array([[x0[0,10],x0[0,11],x0[0,12]],
                            [x0[0,11],x0[0,13],x0[0,14]],
                            [x0[0,12],x0[0,14],x0[0,15]]])
                
                self.v=x0[0,7:10]
                self.kf_P=P
                
                print(np.linalg.det(P))
                
                path_msg = Path()
                path_msg.header.frame_id = 'vrpn'
                path_msg.header.stamp = rospy.Time.now()
                
                pose_array=PoseArray()
                pose_array.header.frame_id="vrpn"
                for i in range(N):
                    pose = PoseStamped()
                    pose.header = path_msg.header
                    pose.pose.position.x = X[i,0]
                    pose.pose.position.y = X[i,1]
                    pose.pose.position.z = X[i,2]
                    pose.pose.orientation.w = 1.0
                    pose.header.frame_id='vrpn' 
                    path_msg.poses.append(pose)
                    
                    T_Bi_V=numpify(pose.pose)
                    
                    v_Bi=transform_vector(np.linalg.inv(T_Bi_V),v_V)
                    roll,pitch,yaw=euler_from_vector(v_Bi)
                    
                    
                    v_desire_Bi=transform_vector(np.linalg.inv(T_Bi_V),v_desire_V)
                    v_target_Bi=transform_vector(np.linalg.inv(T_Bi_V),v_target_V)
                    
                    min_angles,max_angles=calculate_visual_limit(v_target_Bi,fov_pitch=np.pi/2-0.5,fov_yaw=np.pi/2-0.5)
                    
                    
                    roll,pitch,yaw=euler_from_vector(v_desire_Bi)
                    
                    pitch,yaw=np.clip(np.array([pitch,yaw]),min_angles,max_angles)
                    #print(euler_from_vector(v_Bi))
                    T_Bi_V[:3,:3]=euler_matrix(roll,pitch,yaw)[:3,:3]
                    
                    
                    
                    pose=msgify(Pose,T_Bi_V)
                    pose_array.poses.append(pose)
                    
                        

                #print(pose.pose)
                    
                self.path_pub.publish(path_msg)
                self.pose_arr_pub.publish(pose_array)

                rate.sleep()
                
                
            except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):
                rospy.logwarn("Failed to lookup transform!")

            #rate.sleep()

if __name__ == '__main__':
    node = TransformListenerNode()